import ollama
from pymilvus import MilvusClient
from tqdm import tqdm
milvus_client = MilvusClient("http://123.60.131.133:19530")
collection_name = "my_rag_collection"

# 将文本保存到milvus
def save_text_embedding(text_lines: list[str]):
    data = []
    for i, line in enumerate(tqdm(text_lines, desc="Creating embeddings")):
        data.append({"id": i, "vector": emb_text(line), "text": line})

    milvus_client.insert(collection_name=collection_name, data=data)

# 将文本进行emb，转换成数组
def emb_text(text: str):
    return ollama.embed(model='nomic-embed-text', input=text).embeddings[0]

# 初始化数据库链接
def init_collection():
    if milvus_client.has_collection(collection_name):
        milvus_client.drop_collection(collection_name)
    milvus_client.create_collection(
        collection_name=collection_name,
        dimension=768,
        metric_type="IP",  # Inner product distance
        consistency_level="Strong",  # Supported values are (`"Strong"`, `"Session"`, `"Bounded"`, `"Eventually"`). See https://milvus.io/docs/consistency.md#Consistency-Level for more details.
    )

def close():
    milvus_client.close()



if __name__ == '__main__':
    init_collection()
    save_text_embedding([
        "2022年，中国空间站完成在轨建造，空间站工程转入为期10年以上的应用与发展阶段，计划每年发射两艘载人飞船和两艘货运飞船，为满足密集发射任务的需要，天舟货运飞船实行组批生产的方式，天舟六号至天舟十一号是同一批生产的飞船，同一批次的外形、功能相似或相近。 [6]天舟九号货运飞船是空间站应用与发展阶段组批生产的第4艘货运飞船，承担着为神舟二十号和神舟二十一号乘组运送物资的任务，装载了航天员在轨驻留消耗品、推进剂、应用实验装置等物资，共约6.5吨。 [5]2024年10月29日，神舟十九号载人飞行任务新闻发布会通报了天舟九号将于2025年发射的信息 [7]；同日，中国载人航天工程办公室启动天舟九号飞行任务标识征集活动 [8]。2025年1月，天舟九号飞行任务标识正式发布 [9]。7月12日，天舟九号货运飞船和长征七号遥十运载火箭组合体垂直转运至发射区，按计划开展发射前的各项功能检查、联合测试等工作，择机实施发射。 [3] [10]7月13日，天舟九号完成全区合练。"
        ,"天舟九号可为空间站在轨运营和航天员长期驻留提供有力保障和支撑。 [5]飞船发射后将对接于核心舱后向端口，主要上行航天员驻留物资、舱外服等出舱消耗物资，保障平台安全运行的维修备件和推进剂、应用任务各类载荷和样品；下行在轨废弃物。 [1]"
        ])
    close()